#include <pthread.h>
#include <mutex>
#include <thread>
#include <unistd.h>
#include <signal.h>

#include <arpa/inet.h>
#include <sys/ioctl.h>
#include <net/if.h>

#include "cnn_function.hpp"
#ifdef NO_POINT_GREY
#include "opencv2/opencv.hpp"
cv::VideoCapture capture;
#else
#include "pgr_camera.hpp"
#endif
#include "jsmn.h"
#include "util.h"

std::vector<std::string> current_user_name;
std::vector<int> current_user_spoofing;
std::vector<cv::Rect> rois_global;

std::mutex send_image_lock;
bool have_new_send_image = false;
cv::Mat new_send_image;

bool listen_thread_stop = false;
pthread_t stream_thread_pid;

#define PROTO_HEAD_BYTE 4
#define PROTO_TYPE_BYTE 2
#define VERSION_BYTE 4
#define CAMMAND_BYTE 4
#define TOKEN_BYTE 4
#define MSG_HEAD_BYTE (PROTO_HEAD_BYTE + PROTO_TYPE_BYTE + VERSION_BYTE + CAMMAND_BYTE + TOKEN_BYTE)
#define IMAGE_SIZE_BYTE 4

void quit_recognition()
{
    listen_thread_stop = true;
}

void sig_close(int sig){
    const char *s;
    switch (sig) {
    case SIGINT:
        s = "interrupt";
        break;
    case SIGTERM:
        s = "terminate";
        break;
    }
    quit_recognition();
    printf("\n\nreceived signal: %s, quit...\n", s);
}

int little_to_big_endian(int value)
{
    return ((value & 0x000000FF) << 24) | ((value & 0x0000FF00) << 8) | ((value & 0x00FF0000) >> 8) | ((value & 0xFF000000) >> 24);
}

void swap_binary(char *data, int length)
{
    if(length % 2 != 0){
        printf("swap_binary: data length error: %d\n", length);
        exit(-1);
    }
    int half_length = length / 2;
    for(int i = 0; i < half_length; i++){
        char temp = data[i];
        data[i] = data[length - 1 - i];
        data[length - 1 - i] = temp;
    }
}

static int jsoneq(const char *json, jsmntok_t *tok, const char *s)
{
    if (tok->type == JSMN_STRING && (int) strlen(s) == tok->end - tok->start &&
        strncmp(json + tok->start, s, tok->end - tok->start) == 0) {
        return 0;
    }
    return -1;
}

void fill_msg_head(char *msg, short proto_type, int command, int token, int body_length)
{
    int head_length = MSG_HEAD_BYTE - PROTO_HEAD_BYTE + IMAGE_SIZE_BYTE + body_length;
    int version = 1;
    *((int *)msg) = head_length;
    swap_binary(msg, PROTO_HEAD_BYTE);
    *((short *)(msg + PROTO_HEAD_BYTE)) = proto_type;
    swap_binary(msg + PROTO_HEAD_BYTE, PROTO_TYPE_BYTE);
    *((int *)(msg + PROTO_HEAD_BYTE + PROTO_TYPE_BYTE)) = version;
    swap_binary(msg + PROTO_HEAD_BYTE + PROTO_TYPE_BYTE, VERSION_BYTE);
    *((int *)(msg + PROTO_HEAD_BYTE + PROTO_TYPE_BYTE + VERSION_BYTE)) = command;
    swap_binary(msg + PROTO_HEAD_BYTE + PROTO_TYPE_BYTE + VERSION_BYTE, CAMMAND_BYTE);
    *((int *)(msg + PROTO_HEAD_BYTE + PROTO_TYPE_BYTE + VERSION_BYTE + CAMMAND_BYTE)) = token;
    swap_binary(msg + PROTO_HEAD_BYTE + PROTO_TYPE_BYTE + VERSION_BYTE + CAMMAND_BYTE, TOKEN_BYTE);
    *((int *)(msg + PROTO_HEAD_BYTE + PROTO_TYPE_BYTE + VERSION_BYTE + CAMMAND_BYTE + TOKEN_BYTE)) = body_length;
    swap_binary(msg + PROTO_HEAD_BYTE + PROTO_TYPE_BYTE + VERSION_BYTE + CAMMAND_BYTE + TOKEN_BYTE, IMAGE_SIZE_BYTE);
}

bool socket_send_data(int socket_client, char *data, int bytes_left)
{
    int written_bytes = 0;
    bool write_error = false;
    while(bytes_left > 0){
        if(listen_thread_stop) break;
        int write_bytes = write(socket_client, data + written_bytes, bytes_left);
        if(write_bytes < 0){
            write_error = true;
            printf("write error: write return: %d, %d socket_client: %d\n", write_bytes, bytes_left, socket_client);
            break;
        } else if(write_bytes == 0) {
            printf("write error: write return: %d, %d socket_client: %d\n", write_bytes, bytes_left, socket_client);
            break;
        } else {
            bytes_left -= write_bytes;
            written_bytes += write_bytes;
            printf("write return: %d, bytes_left %d written_bytes %d socket_client: %d\n",
                   write_bytes, bytes_left, written_bytes, socket_client);
        }
    }
    return write_error;
}

int send_json_data(int socket_client, char *msg, int length, int *send_msg_num_ret)
{
    const static unsigned int msg_max_size = 64 * 1024 + IMAGE_SIZE_BYTE + MSG_HEAD_BYTE;
    static int send_msg_num = 10000 * 10000;
    static char msg_data[msg_max_size];
    if(length <= msg_max_size - IMAGE_SIZE_BYTE - MSG_HEAD_BYTE){
        memcpy(msg_data + IMAGE_SIZE_BYTE + MSG_HEAD_BYTE, msg, length * sizeof(char));
        fill_msg_head(msg_data, 1, 10, send_msg_num, length);
    } else {
        printf("send_json_data failed: msg too large!");
        return 1;
    }
    bool write_error = socket_send_data(socket_client, msg_data, length + IMAGE_SIZE_BYTE + MSG_HEAD_BYTE);
    if(write_error){
        quit_recognition();
    } else {
        *send_msg_num_ret = send_msg_num;
        send_msg_num += 1;
    }
    return 0;
}

int send_image_data(int socket_client, cv::Mat &img, int *send_msg_num_ret)
{
    const static unsigned int image_max_size = 2 * 1024 * 1024 + IMAGE_SIZE_BYTE + MSG_HEAD_BYTE;
    static char img_encode_buf[image_max_size];
    static int send_image_num = 0;
    std::vector<uchar> img_encode;

    cv::imencode(".jpg", img, img_encode);
    if(img_encode.size() <= image_max_size - IMAGE_SIZE_BYTE - MSG_HEAD_BYTE){
        memcpy(img_encode_buf + IMAGE_SIZE_BYTE + MSG_HEAD_BYTE, img_encode.data(), img_encode.size() * sizeof(char));
        fill_msg_head(img_encode_buf, 2, 2001, send_image_num, img_encode.size());
    } else {
        printf("send_image_data failed: image too large!");
        return 1;
    }
    printf("send_image_data: image size %lu, total size: %lu\n", img_encode.size(), img_encode.size() + IMAGE_SIZE_BYTE + MSG_HEAD_BYTE);
    bool write_error = socket_send_data(socket_client, img_encode_buf, img_encode.size() + IMAGE_SIZE_BYTE + MSG_HEAD_BYTE);
    if(write_error){
        quit_recognition();
    } else {
        *send_msg_num_ret = send_image_num;
        send_image_num += 1;
    }
    return 0;
}

char *receive_server_msg(int socket_client, int *msg_length)
{
    static char buf[16 * 1024];
    char *server_msg = 0;
    int bytes_left = 1;
    int bytes_has_read = 0;
    while(bytes_left > 0){
        int bytes_read = read(socket_client, buf, sizeof(buf));
        if(bytes_read <= 0) {
            printf("read error: return: %d\n", bytes_read);
            quit_recognition();
            break;
        } else {
            if(bytes_has_read == 0){
                if(bytes_read < 4){
                    quit_recognition();
                    break;
                } else {
                    swap_binary(buf, PROTO_HEAD_BYTE);
                    *msg_length = *((int *)buf);
                    server_msg = (char *)calloc(*msg_length + 1, sizeof(char));
                    memcpy(server_msg, buf + PROTO_HEAD_BYTE, bytes_read - PROTO_HEAD_BYTE);
                }
                bytes_left = *msg_length - bytes_read + 4;
                bytes_has_read = bytes_read - 4;
            } else {
                memcpy(server_msg + bytes_has_read, buf, bytes_read);
                bytes_left -= bytes_read;
                bytes_has_read += bytes_read;
            }
            printf("msg_length: %d, bytes_left: %d, bytes_has_read: %d, read return: %d\n",
                   *msg_length, bytes_left, bytes_has_read, bytes_read);
        }
    }
    //printf("server_msg: %s\n", server_msg);
    return server_msg;
}

int parse_server_msg(char *msg, int msg_length, int *msg_index)
{
    swap_binary(msg, PROTO_TYPE_BYTE);
    short proto_type = *((short *)msg);
    swap_binary(msg + PROTO_TYPE_BYTE, VERSION_BYTE);
    int version = *((int *)(msg + PROTO_TYPE_BYTE));
    swap_binary(msg + PROTO_TYPE_BYTE + VERSION_BYTE, CAMMAND_BYTE);
    int command = *((int *)(msg + PROTO_TYPE_BYTE + VERSION_BYTE));
    swap_binary(msg + PROTO_TYPE_BYTE + VERSION_BYTE + CAMMAND_BYTE, TOKEN_BYTE);
    int token = *((int *)(msg + PROTO_TYPE_BYTE + VERSION_BYTE + CAMMAND_BYTE));
    swap_binary(msg + PROTO_TYPE_BYTE + VERSION_BYTE + CAMMAND_BYTE + TOKEN_BYTE, IMAGE_SIZE_BYTE);
    int msg_data_length = *((int *)(msg + PROTO_TYPE_BYTE + VERSION_BYTE + CAMMAND_BYTE + TOKEN_BYTE));
    if(msg_length - (PROTO_TYPE_BYTE + VERSION_BYTE + CAMMAND_BYTE + TOKEN_BYTE + IMAGE_SIZE_BYTE) != msg_data_length){
        printf("parse_server_msg message data length error\n");
        quit_recognition();
        return 0;
    }
    printf("proto_type %d, version %d, command %d, token %d, msg_data_length %d,\n",
           proto_type, version, command, token, msg_data_length);

    char *msg_content = msg + PROTO_TYPE_BYTE + VERSION_BYTE + CAMMAND_BYTE + TOKEN_BYTE + IMAGE_SIZE_BYTE;
    jsmn_parser msg_parser;
    jsmn_init(&msg_parser);
    int token_num = jsmn_parse(&msg_parser, msg_content, strlen(msg_content), NULL, 0);
    if(token_num < 1){
        printf("parse_server_msg: Failed to parse JSON: token_num %d, msg_content: %s\n", token_num, msg_content);
        quit_recognition();
        return 0;
    }
    jsmntok_t *msg_token = (jsmntok_t *)calloc(token_num, sizeof(jsmntok_t));
    jsmn_init(&msg_parser);
    token_num = jsmn_parse(&msg_parser, msg_content, strlen(msg_content), msg_token, token_num);
    if(token_num < 1){
        printf("parse_server_msg: 1 Failed to parse JSON: token_num %d, msg_content: %s\n", token_num, msg_content);
        quit_recognition();
        return 0;
    }
    printf("msg_content: %s token_num: %d\n", msg_content, token_num);
    int return_code = 0;
    for(int i = 1; i < token_num; i++) {
        if (jsoneq(msg_content, &msg_token[i], "responseCode") == 0) {
            printf("responseCode: %.*s\n", msg_token[i+1].end - msg_token[i+1].start, msg_content + msg_token[i+1].start);
            return_code = atoi(msg_content + msg_token[i+1].start);
            i++;
        } else {
            printf("unexpected key: %.*s\n", msg_token[i].end - msg_token[i].start, msg_content + msg_token[i].start);
        }
    }
    *msg_index = token;
    free(msg_token);
    return return_code;
}

void *listen_thread(void *arg)
{
    int socket_desc = socket(AF_INET, SOCK_STREAM , 0);
    if (socket_desc == -1) {
        printf("Could not create socket\n");
        quit_recognition();
        return 0;
    }
    int opt = 1;
    setsockopt(socket_desc, SOL_SOCKET, SO_REUSEADDR, (const void *)&opt, sizeof(opt));
    struct timeval time_out_send = {6, 0};
    socklen_t time_out_length = sizeof(time_out_send);
    setsockopt(socket_desc, SOL_SOCKET, SO_SNDTIMEO, &time_out_send, time_out_length);
    struct timeval time_out_rec = {10, 0};
    time_out_length = sizeof(time_out_rec);
    setsockopt(socket_desc, SOL_SOCKET, SO_RCVTIMEO, &time_out_rec, time_out_length);

    struct sockaddr_in client;
    client.sin_family = AF_INET;
    client.sin_addr.s_addr = inet_addr("192.168.3.181");
    client.sin_port = htons(10011);
    if(connect(socket_desc, (struct sockaddr *)&client, sizeof(client))==-1){
        perror("client connect failed\n");
        quit_recognition();
        return 0;
        //exit(-1);
    } else {
        printf("client connect success\n");
    }
    char msg[256] = {0};
    char client_id[13] = {0};
    getMacAddress(client_id);
    sprintf(msg, "{\"clientType\":3000,\"clientId\":\"%s\"}", client_id);
    printf("json : %s\n", msg);
    int send_msg_num_ret = 0;
    send_json_data(socket_desc, msg, strlen(msg), &send_msg_num_ret);
    int msg_length = 0;
    char *server_msg = receive_server_msg(socket_desc, &msg_length);
    if(server_msg == 0){
        quit_recognition();
        return 0;
    }
    int msg_index = 0;
    int return_code = parse_server_msg(server_msg, msg_length, &msg_index);
    if(send_msg_num_ret != msg_index || return_code != 200){
        printf("error: login error: %d %d return_code: %d\n", send_msg_num_ret, msg_index, return_code);
        quit_recognition();
    }
    free(server_msg);
    printf("login success\n");

    int send_frame = 0;
    double send_time = getMsecOfNow();
    while (true) {
        if(listen_thread_stop) break;
        cv::Mat send_image_clone;
        while (true) {
            if(listen_thread_stop) break;
            {
                std::lock_guard<std::mutex> send_image_lock_ws(send_image_lock, std::adopt_lock);
                if(have_new_send_image){
                    send_image_clone = new_send_image.clone();
                    have_new_send_image = false;
                    break;
                }
            }
            usleep(30000);
        }
        printf("send_image_clone copy over\n");
        if(send_image_clone.rows <= 0 || send_image_clone.cols <= 0 || listen_thread_stop) break;
        send_image_data(socket_desc, send_image_clone, &send_msg_num_ret);
        server_msg = receive_server_msg(socket_desc, &msg_length);
        if(server_msg == 0){
            quit_recognition();
            return 0;
        }
        msg_index = 0;
        return_code = parse_server_msg(server_msg, msg_length, &msg_index);
        if(send_msg_num_ret != msg_index){
            printf("error: msg token id diff: %d %d return_code: %d\n", send_msg_num_ret, msg_index, return_code);
            quit_recognition();
        }
        free(server_msg);
        if(return_code == 400) printf("\n\nrecognition success\n\n");
        else printf("\n\nrecognition failed\n\n");
        double now_time = getMsecOfNow();
        if(now_time - send_time > 5000){
            send_frame += 1;
            printf("fps: %f\n", send_frame * 1000 / (now_time - send_time));
            send_frame = 0;
            send_time = now_time;
        } else {
            send_frame += 1;
        }

        //printf("try 1 round finished\n");;
        //exit(-1);
    }
    close(socket_desc);
    printf("listen_thread quit...\n");
}

void copy_mat_data(cv::Mat &src, cv::Mat &dst, int width_offset, int height_offset)
{
    int channels = src.channels();
    int face_height = src.rows;
    int face_width = src.cols;
    for(int k= 0; k < channels; ++k){
        for(int m = 0; m < face_height; ++m){
            for(int n = 0; n < face_width; ++n){
                dst.at<cv::Vec3b>(m + height_offset, n + width_offset)[k] = src.at<cv::Vec3b>(m, n)[k];
            }
        }
    }
}
int main(int argc, char **argv)
{
    signal(SIGINT, sig_close);  //close CTRL+C
    signal(SIGTERM, sig_close);
#ifdef NO_POINT_GREY
    capture.open(0);
    if (!capture.isOpened()) {
        printf("Video capture not successfully initialized, exit!\n");
        //return 0;
    }
    capture.set(CV_CAP_PROP_FRAME_WIDTH, 480);
    capture.set(CV_CAP_PROP_FRAME_HEIGHT, 640);
    capture.set(CV_CAP_PROP_FPS, 10);
#else
    PGRCam cam;
    cam.initCamera();
#endif
    //init_detector("/var/darknet/cnn/cfg/yolov3-spp.cfg", "/var/darknet/cnn/yolov3-spp_final.weights");
    //init_detector("/var/darknet/cnn/cfg/yolov3-tiny.cfg", "/var/darknet/cnn/yolov3-tiny_final.weights");
    init_detector("/var/darknet/cnn/cfg/yolov3-small.cfg", "/var/darknet/cnn/yolov3-small_final.weights");
    //init_detector("/data/aaa/yolov3-small.cfg", "/data/aaa/yolov3-small_final.weights");
    //init_detector("/data/aaa/yolov3-spp.cfg", "/data/aaa/yolov3-spp_final.weights");
    int detection_bbox[MAX_BBOX_NUM * 4];
    pthread_create(&stream_thread_pid, NULL, listen_thread, NULL);
    cv::Mat image_input;

    std::vector<cv::Mat> aligned_faces;
    while(true){
        if(listen_thread_stop) break;
#ifdef NO_POINT_GREY
        bool bCap = true;
        bCap = capture.read(image_input);
        //image_input = cv::imread("/sdcard/A/test.png");
#else
        bool bCap = cam.grabFrame();
#endif
        if (!bCap) {
            printf("warning: frame not captured\n");
            usleep(20000);
            continue;
        }
#ifndef NO_POINT_GREY
        image_input = cam.rgb_image_cv;
#endif
        
        int face_count = 0;
        detection_face(image_input, &face_count, detection_bbox);
        printf("detection_face over face_count %d\n", face_count);
        if(face_count == 0) continue;

        aligned_faces.clear();
        aligned_faces.resize(face_count);
        std::vector<cv::Size> send_image_size;
        int send_image_width = 0;
        int send_image_height = 0;
        int face_size_max = 256;
        for(int i = 0; i < face_count; i++){
            cv::Rect roi_local(cv::Point(detection_bbox[i * 4], detection_bbox[i * 4 + 1]),
                               cv::Size(detection_bbox[i * 4 + 2] - detection_bbox[i * 4],
                                        detection_bbox[i * 4 + 3] - detection_bbox[i * 4 + 1]));
            //cv::rectangle(image_input, cvPoint(roi_local.x, roi_local.y),
            //              cvPoint(roi_local.x + roi_local.width, roi_local.y + roi_local.height), cvScalar(255,0,0), 4);
            cv::Rect roi_real;
            roi_real.x = std::max(roi_local.x - int(roi_local.width * 0.1), 0);
            roi_real.y = std::max(roi_local.y - int(roi_local.height * 0.1), 0);
            roi_real.width = std::min(int(roi_local.width * 1.2), image_input.cols - roi_real.x - 1);
            roi_real.height = std::min(int(roi_local.height * 1.2), image_input.rows - roi_real.y - 1);
            int local_width = 0;
            int local_height = 0;
            cv::Mat face_resized;
            if(roi_real.width > face_size_max || roi_real.height > face_size_max){
                int new_w = roi_real.width;
                int new_h = roi_real.height;
                if (((float)face_size_max / roi_real.width) < ((float)face_size_max / roi_real.height)) {
                    new_w = face_size_max;
                    new_h = (roi_real.height * face_size_max) / roi_real.width;
                } else {
                    new_h = face_size_max;
                    new_w = (roi_real.width * face_size_max) / roi_real.height;
                }

                cv::resize(image_input(roi_real), face_resized, cv::Size(new_w, new_h));
                aligned_faces[i]  = face_resized;
                local_width = new_w;
                local_height = new_h;
            } else {
                aligned_faces[i]  = image_input(roi_real);
                local_width = roi_real.width;
                local_height = roi_real.height;
            }
            send_image_size.push_back(cv::Size(local_width, local_height));
            if(i % 2 == 0){
                send_image_width += local_width;
                send_image_height += local_height;
            } else {
                send_image_width += local_width;
                if(local_height > send_image_size[i - 1].height) send_image_height += (local_height - send_image_size[i - 1].height);
            }
        }
        cv::Mat send_image_mat(send_image_height, send_image_width, CV_8UC3, cv::Scalar(0));
        int send_image_width_offset = 0;
        int send_image_height_offset = 0;
        for(int i = 0; i < face_count; i++){
            if(i % 2 == 0){
                copy_mat_data(aligned_faces[i], send_image_mat, send_image_width_offset, send_image_height_offset);
                send_image_width_offset += send_image_size[i].width;
            } else {
                copy_mat_data(aligned_faces[i], send_image_mat, send_image_width_offset, send_image_height_offset);
                send_image_width_offset = 0;
                send_image_height_offset += std::max(send_image_size[i].height, send_image_size[i - 1].height);
            }
        }
        //cv::imwrite("ttt.jpg", send_image_mat);
        //cv::imwrite("grab.jpg", image_input);
        {
            std::lock_guard<std::mutex> send_image_lock_ws(send_image_lock, std::adopt_lock);
            new_send_image = send_image_mat.clone();
            have_new_send_image = true;
        }
        //if(face_count != 0) break;
        //usleep(1000000);
    }
    pthread_join(stream_thread_pid, NULL);
    uninit_detector();
    printf("main quit...\n");
    return 0;
}
